# Alexander F. Gazmararian
# afg2@princeton.edu

library(tidyverse)
library(here)
library(scales)
library(modelsummary)
library(kableExtra)
library(margins)
library(broom)
library(lmtest)
library(sandwich)
library(gridExtra)
library(viridis)
library(ggsci)
library(emmeans)
library(marginaleffects)
library(MASS)

# Load custom functions
source(here("code", "fun", "fix_txt.r"))
source(here("code", "fun", "table_spec.r"))

# Load data
nat <- readRDS(here("data", "national_wide.rds"))
fair <- readRDS(here("data", "fair_wide.rds"))

# Plot set up
p_dodge <- .91

# Prepare data
fair$support <- ordered(fair$support, levels = c("Very unlikely", "Somewhat unlikely", "Somewhat likely", "Very likely"))

nat$sample <- "National"
fair$sample <- "Fossil Fuel"

df.all <- rbind(subset(nat, select = c(support, weights, sample), transition == 1), fair[, c("support", "weights", "sample")])


# Create Figure 3 ----

## Fossil fuel employment ----

p.ff <-
  fair %>%
  group_by(support, sample, ffemploy) %>%
  summarise(n = sum(weights)) %>%
  group_by(sample, ffemploy) %>%
  mutate(pct = n / sum(n),
         ffemploy = ifelse(ffemploy == 1, "Fossil Fuels", "Other Industry")) %>%
  ggplot(aes(x=support,y=pct,group=ffemploy,fill=ffemploy, label=percent(pct, accuracy = 1))) +
  geom_col(position=position_dodge(p_dodge)) +
  geom_text(position=position_dodge(p_dodge), vjust = -.5) +
  theme_bw(base_size = 14) +
  scale_y_continuous(limits = c(0, .75), breaks = seq(0, .75, .25), labels = percent, expand = c(0,0)) +
  scale_x_discrete(labels = label_wrap(15)) +
  labs(x = "",
       y = "Percent (Weighted)",
       fill = "",
       title = expression(bold("B") ~ " Fossil Fuel Employment")) +
  scale_fill_manual(values = as.vector(palette.colors(palette = "Okabe-Ito")[c(4,7)])) +
  theme(
    legend.position = c(.175, .9),
    panel.grid = element_blank(),
    legend.background = element_blank(),
    plot.margin = margin(b = -5, l = 5)
  )
p.ff

## Party identification ----
p.pid <-
  fair %>%
  group_by(support, partysummary) %>%
  summarise(n = sum(weights)) %>%
  group_by(partysummary) %>%
  mutate(pct = n / sum(n)) %>%
  ggplot(aes(x=support,y=pct,group=partysummary,fill=partysummary, label=percent(pct, accuracy = 1))) +
  geom_col(position=position_dodge(p_dodge)) +
  geom_text(position=position_dodge(p_dodge), vjust = -.5) +
  theme_bw(base_size = 14) +
  scale_y_continuous(limits = c(0, .75), breaks = seq(0, .75, .25), labels = percent, expand = c(0,0)) +
  scale_x_discrete(labels = label_wrap(15)) +
  scale_fill_manual(values = as.vector(palette.colors(palette = "Okabe-Ito", n = 3))) +
  labs(x = "",
       y = "Percent (Weighted)",
       fill = "",
       title = expression(bold("C") ~ " Partisan Identification")) +
  theme(
    legend.position = c(.15, .85),
    panel.grid = element_blank(),
    legend.background = element_blank(),
    plot.margin = margin(b = -5, l = 5)
  )
p.pid

## Community versus national ----
p.support <-
  df.all %>%
  group_by(support, sample) %>%
  summarise(n = sum(weights)) %>%
  group_by(sample) %>%
  mutate(
    pct = n / sum(n)
  ) %>%
  ggplot(aes(x=support,y=pct,group=sample,fill=sample, label=percent(pct, accuracy = 1))) +
  geom_col(position=position_dodge(p_dodge)) +
  geom_text(position=position_dodge(p_dodge), vjust = -.5) +
  theme_bw(base_size = 14) +
  scale_y_continuous(limits = c(0, .75), breaks = seq(0, .75, .25), labels = percent, expand = c(0,0)) +
  scale_x_discrete(labels = label_wrap(15)) +
  labs(x = "",
       y = "Percent (Weighted)",
       fill = "",
       title = expression(bold("A") ~ " Fossil Fuel Communities vs. National Public")) +
  scale_fill_manual(values = as.vector(palette.colors(palette = "Okabe-Ito")[c(6,10)])) +
  theme(
    legend.position = c(.15, .85),
    panel.grid = element_blank(),
    legend.background = element_blank(),
    plot.margin = margin(b = -5, l = 5)
  )
p.support

## Estimate effect of assistance ----

#Coefficient map
coefnames <- c(
  "transition"="Treatment: Transition Assistance",
  "transition:partysummaryNeither"="Treatment x Independent",
  "transition:partysummaryRepublican"="Treatment x Republican",
  "age"="Age",
  "sexMale"="Male",
  "black"="Black",
  "hispanic"="Hispanic",
  "ff"="Fossil Fuel Employment",
  "fullemploy"="Employed",
  "college"="College Degree",
  "income5Less than $29,999"="Income: 1st Quartile",
  "income5$30,000 - $59,999"="Income: 2nd Quartile",
  "income5$60,000 - $99,999"="Income: 3rd Quartile",
  "income5Prefer not to say"="Income: Not Say",
  "partysummaryRepublican" = "Republican",
  "partysummaryNeither" = "Independent",
  "(Intercept)" = "Intercept"
)

m <- list()
f.base <- y ~ partysummary + age + sex + college + black + hispanic + ff + fullemploy + income5
m[[1]] <- lm(support_binary ~ transition, nat)
m[[2]] <- lm(update(f.base, support_binary ~ transition + .), nat)
m[[3]] <- lm(update(f.base, support_binary ~ transition + .), nat, weights = weights)
m[[4]] <- lm(update(f.base, support_binary ~ transition * partysummary + .), nat)
m[[5]] <- lm(update(f.base, as.numeric(support) ~ transition), nat)
m[[6]] <- lm(update(f.base, as.numeric(support) ~ transition + .), nat)

### Create SI Table 11.1 ----
file_trans <- here("output", "tables", "si_tab_11.1.txt")
modelsummary(
  m,
  vcov = "HC3",
  stars = c("*"=.1,"**"=.05,"***"=.01),
  coef_map = coefnames,
  gof_map = gm,
  output = "latex",
  escape = FALSE,
  add_rows = data.frame(
    "Sample Weights",
    "No",
    "No",
    "Yes",
    "No",
    "No",
    "No"
  )
) %>%
  kableExtra::add_header_above(c(" " = 1, "Binary" = 4, "Scale" = 2)) %>%
  cat(., file = file_trans)
fix_txt(file_trans)

# Create plot of treatment effect
ame <- tidy(coeftest(m[[2]], vcovHC(m[[2]], "HC2")), conf.int = TRUE)
ame <- ame[2, ]
ame

# Create plot of heterogeneous treatment effects

### Create ATE Plot ----
ate1 <- tidy(coeftest(m[[1]], vcovHC(m[[1]], "HC2")), conf.int = TRUE)
ate1$Model <- "No Controls\n"
ate2 <- tidy(coeftest(m[[2]], vcovHC(m[[2]], "HC2")), conf.int = TRUE)
ate2$Model <- "Covariate Adjusted\n"

ate <- rbind(ate1, ate2)

ate <- subset(ate, term == "transition")

ate$lb90 <- with(ate, estimate + std.error * qnorm(.05))
ate$ub90 <- with(ate, estimate + std.error * qnorm(.95))

p.ate <-
  ate %>%
  ggplot(aes(x = Model, y = estimate)) +
  geom_hline(yintercept = 0, color = "grey") +
  geom_errorbar(aes(ymin = conf.low, ymax = conf.high), width = 0) +
  geom_errorbar(aes(ymin = lb90, ymax = ub90), width = 0, linewidth = 1.25) +
  geom_point(size = 3, fill = "white", shape = 21) +
  theme_bw(base_size = 14) +
  scale_y_continuous(limits = c(-.06, .16), labels = percent, expand = c(0,0)) +
  theme(
    panel.grid = element_blank(),
    legend.position = "",
    plot.margin = margin(b = -5, l = 5)
  ) +
  labs(y = "Change in Climate Policy Support",
       x = "",
       title = expression(bold("D") ~ " Average Treatment Effect (National)"))
p.ate


p.all <-
  nat %>%
  group_by(support_binary, transition) %>%
  summarize(n = sum(weights)) %>%
  group_by(transition) %>%
  mutate(pct = n / sum(n)) %>%
  mutate(transition = ifelse(transition == 1, "Treatment", "Control"),
         support_binary = ifelse(support_binary == 1, "Likely\n", "Unlikely\n")) %>%
  ggplot(aes(x = support_binary, y = pct, fill = transition, label=percent(pct, accuracy = 1))) +
  geom_col(position = position_dodge(p_dodge)) +
  geom_text(position = position_dodge(p_dodge), vjust = -.5) +
  geom_text(data = data.frame(
    support_binary = c("Unlikely\n", "Likely\n"),
    transition = c("Treatment", "Control"),
    pct = c(.3, .9),
    label = c("***", "***")
  ),
  aes(label = label),
  position = position_dodge(p_dodge)
  ) +
  theme_bw(base_size = 14) +
  scale_fill_manual(values = c("darkgrey", pal_npg("nrc")(10)[[4]])) +
  labs(x = "Climate Policy Support Likelihood", y = "Percent", 
       title = expression(bold("C") ~ " Causal Effect of Transition Assistance"),
       fill = "") +
  scale_y_continuous(limits = c(0, 1.05), labels = percent, expand = c(0,0)) +
  scale_x_discrete(limits = rev) +
  theme(
    legend.position = c(.15, .95),
    panel.grid = element_blank(),
    legend.background = element_blank()
  )
p.all


emean.pid <- emmeans(m[[4]], specs = ~ partysummary * transition)

p.pid <-
  emean.pid %>%
  data.frame() %>%
  mutate(transition = ifelse(transition == 1, "Treatment", "Control")) %>%
  mutate(partysummary = paste0(partysummary, "\n")) %>%
  ggplot(aes(x = partysummary, y = emmean, fill = transition, label = percent(emmean, accuracy = 1))) +
  geom_col(position = position_dodge(p_dodge)) +
  geom_text(position = position_dodge(p_dodge), vjust = -.5) +
  geom_text(
    data = data.frame(
      partysummary = c("Neither\n", "Republican\n", "Democrat\n"),
      label = c("***", "**", ""),
      transition = rep("Treatment", 3),
      emmean = c(.88, .72, 1.02)
    ),
    aes(label = label),
    position = position_dodge(p_dodge)
  ) +
  theme_bw(base_size = 14) +
  scale_fill_manual(values = c("darkgrey", pal_npg("nrc")(10)[[4]])) +
  scale_y_continuous(limits = c(0,1.05), expand = c(0,0), labels = percent) +
  labs(x = "Partisan Identification",
       title = expression(bold("C") ~ " Heterogeneous Effects by Party"),
       y = "Climate Policy Support", fill = "") +
  theme(
    panel.grid = element_blank(),
    legend.position = c(.85, .95),
    legend.background = element_blank()
  )
p.pid

## Combine plots to make Figure 3 ----
p.out <- grid.arrange(p.support, p.ff, p.pid, p.ate, ncol = 2, nrow = 2)
ggsave(
  p.out, 
  filename = here("output", "figures", "fig_3.png"),
  scale = 1.75,
  dpi = 300,
  width = 6.5,
  height = 4
)

# Correlates of support: SI Table 10.3 ----
m.olr <-
  MASS::polr(
    formula = support ~ age_tri + sex + college + ffemploy + partysummary + income_acs + fair + diversify + gw_binary,
    data = fair,
    weights = weights
  )
m.olr.uw <-
  MASS::polr(
    formula = support ~ age_tri + sex + college + ffemploy + partysummary + income_acs + fair + diversify + gw_binary,
    data = fair
  )
m.olr.welfare <-
  MASS::polr(
    formula = support ~ age_tri + sex + college + ffemploy + partysummary + income_acs + fair + diversify + gw_binary + welfare_sup,
    data = fair,
    weights = weights
  )
m.ols <-
  lm(
    as.numeric(support) ~ age_tri + sex + college + ffemploy + partysummary + income_acs + fair + diversify + gw_binary,
    data = fair,
    weights = weights
  )
m.ols.uw <-
  lm(
    as.numeric(support) ~ age_tri  + sex + college + ffemploy + partysummary + income_acs + fair + diversify + gw_binary,
    data = fair
  )
m.ols.welfare <- 
  lm(
    as.numeric(support) ~ age_tri  + sex + college + ffemploy + partysummary + income_acs + fair + diversify + gw_binary + welfare_sup,
    data = fair,
    weights = weights
  )
models <- list(
  "(1)" = m.olr.uw,
  "(2)" = m.olr,
  #  "(3)" = m.olr.welfare,
  "(3)" = m.ols.uw,
  "(4)" = m.ols
  # "(6)" = m.ols.welfare
)
rows <- tribble(~term, ~`(1)`, ~`(2)`, ~`(3)`, ~`(4)`,
                "Sample Weights", "No", "Yes", "No", "Yes", 
                "Fair Fixed Effects", "Yes", "Yes", "Yes", "Yes")
attr(rows, 'position') <- c(35,36)


modelsummary(
  models,
  stars = c("*"=.1,"**"=.05,"***"=.01),
  coef_map = c(
    "Very unlikely|Somewhat unlikely"="Very unlikely/Somewhat unlikely",
    "Somewhat unlikely|Somewhat likely"="Somewhat unlikely/Somewhat likely",
    "Somewhat likely|Very likely"="Somewhat likely/Very likely",
    "(Intercept)" = "Intercept",
    "age_tri35-54" = "35-54 years",
    "age_tri55 or older" = "$>$55 years",
    "sex"="Female",
    "white" = "White",
    "college"="College Degree",
    "ffemploy"="Fossil Fuel Employment",
    #"fulltime"="Employed",
    "income_acsLess than $20,000" = "$<$\\$20,000",
    "income_acs$20,000 to $39,999" = "\\$20,000-39,999",
    "income_acs$40,000 to $59,999" = "\\$40,000-59,999",
    "income_acs$60,000 - $99,999" = "\\$60,00-99,999",
    "partysummaryRepublican" = "Republican",
    "partysummaryNeither" = "Independent",
    "gw_binary" = "Climate Concern",
    #"welfare_supOppose" = "Economic Conservatism",
    "diversify" = "Diversify Treatment"
  ),
  vcov = list(NULL, NULL, "HC3", "HC3"),
  gof_map = gm,
  gof_omit = "R2|F|RMSE|Std|AIC|Log.Lik",
  add_rows = rows,
  caption = "Regression models of the determinants of climate policy support in exchange for transition assistance \\label{tab:regression_support}",
  escape = FALSE,
  output = "latex"
) %>%
  kableExtra::kable_styling(position = "center", font_size = 10.25) %>%
  kableExtra::add_header_above(c("", "Ordered Logit" = 2, "Linear" = 2)) %>%
  kableExtra::pack_rows("Intercepts", 1, 8, bold = FALSE) %>%
  kableExtra::pack_rows("Age (Baseline: 18-34 years)", 9, 12, bold = FALSE) %>%
  kableExtra::pack_rows("Income (Baseline: $>$\\$100,000)", 19, 25, bold = FALSE, escape = FALSE) %>%
  kableExtra::pack_rows("Party (Baseline: Democrat)", 27, 30, bold = FALSE) %>%
  kableExtra::add_footnote(
    paste(
      "Notes:",
      "HC3 standard errors employed in the linear regression model.",
      "Less than 2.5% of missing income and sex observations imputed with median response."
    ),
    threeparttable = TRUE,
    notation = "none"
  ) %>%
  cat(., file = here("output", "tables", "si_tab_10.1.txt"))

